import torch
from aio import EncapsulatedNTM

# 输入NTM的特征的大小
sequence_width = 4096
# 经过Controller模块后提取的特征的大小
controller_size = 1024
# Controller模块中具体实现层数，比如：如果Controller使用LSTM实现，那就代表使用1层LSTM网络
controller_layers = 1
# 读写头的数量
num_heads = 1
# 记忆单元Memory的大小，可以理解为对应输入的步长
memory_n = 40
# 记忆单元Memory中每个特征的大小，可以理解为经过Controller特征提取后的大小
memory_m = 1024

# 初始化NTM网络
net = EncapsulatedNTM(sequence_width, sequence_width,
                      controller_size, controller_layers,
                      num_heads, memory_n, memory_m)


# 假设数据的batch_size为50
batch_size = 50
# 假设我们要输入数据序列长度为40，例如视频40帧、一句话40个单词
ntm_steps = 40
# 假设数据的特征维度为4096，例如视频一帧经过特征提取后是4096维度的向量
dim_feature = 4096

# 初始化数据
batch_data = torch.randn(size=[batch_size, ntm_steps, dim_feature])

# 重要：每一次进行batch_size数据训练的时候，要先把NTM的状态置为初始状态，
# 否则本次batch_size训练会带有上一次batch_size数据的信息
net.init_sequence(50)

# 先将batch_size数据都给NTM
for i in range(40):
    net(batch_data[:, i, :])

# 再逐个读取，这样每一个输出都会包含之前所有输入的信息
result = torch.zeros(batch_data.size())
for i in range(40):
    result[:, i,:], _ = net()

print(result.shape)